import tiktoken
from app.client.MinioClient import minio_client
from app.client.OpenaiClient import openai_client
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.text_splitter import RecursiveCharacterTextSplitter, Language

PROMPT_1 = '''
#01 你是一个问答对数据集处理专家。
#02 你的任务是根据我给出的内容，生成适合作为问答对数据集的问题。
#03 问题要尽量短，不要太长。
#04 一句话中只能有一个问题。
#05 生成的问题必须宏观、价值，不要生成特别细节的问题。
#06 生成问题示例：
"""
权益型基金的特点有哪些方面？
介绍一下产品经理。
"""
#07 以下是我给出的内容：
"""
{{此处替换成你的内容}}
"""
'''
PROMPT_2 = '''
#01 你是一个心理学问答对数据集处理专家。
#02 你的任务是根据我的问题和我给出的内容，生成对应的问答对。
#03 站在心理学的角度，答案要全面，多使用我的信息，内容要更丰富。
#04 你必须根据我的问答对示例格式来生成：
"""
{"content": "基金分类有哪些", "summary": "根据不同标准，可以将证券投资基金划分为不同的种类：（1）根据基金单位是否可增加或赎回，可分为开放式基金和封闭式基金。开放式基金不上市交易（这要看情况），通过银行、券商、基金公司申购和赎回，基金规模不固定；封闭式基金有固定的存续期，一般在证券交易场所上市交易，投资者通过二级市场买卖基金单位。（2）根据组织形态的不同，可分为公司型基金和契约型基金。基金通过发行基金股份成立投资基金公司的形式设立，通常称为公司型基金；由基金管理人、基金托管人和投资人三方通过基金契约设立，通常称为契约型基金。我国的证券投资基金均为契约型基金。（3）根据投资风险与收益的不同，可分为成长型、收入型和平衡型基金。（4）根据投资对象的不同，可分为股票基金、债券基金、货币基金和混合型基金四大类。"}
{"content": "基金是什么", "summary": "基金，英文是fund，广义是指为了某种目的而设立的具有一定数量的资金。主要包括公积金、信托投资基金、保险基金、退休基金，各种基金会的基金。从会计角度透析，基金是一个狭义的概念，意指具有特定目的和用途的资金。我们提到的基金主要是指证券投资基金。"}
#05 我的问题如下：
"""
{{此处替换成你上一步生成的问题}}
"""
#06 我的内容如下：
"""
{{此处替换成你的内容}}
"""
'''


class GenqaManger:
    def __init__(self) -> None:
        self.client = openai_client
        self.enc = tiktoken.get_encoding("cl100k_base")

    def __generate_chat_completion(
        self, messages, model="gpt-3.5-turbo-16k", temperature=0, top_p=1
    ):
        """
        通过给定的client，使用指定的模型和参数生成聊天完成内容。

        :param client: API客户端对象，用于发起请求。
        :param messages: 用户的输入或问题，作为生成的起始内容。
        :param model: 要使用的模型名称，默认为"gpt-3.5-turbo-16k"。
        :param temperature: 采样温度，控制输出的随机性，0表示最小随机性，默认为0。
        :param top_p: 核采样中的最高概率，用来截断候选词的概率分布，默认为1。
        :return: 从API得到的聊天完成响应对象。
        """
        try:
            response = self.client.chat.completions.create(
                messages=messages,
                model=model,
                temperature=temperature,
                top_p=top_p,
            )
            return response
        except Exception as e:
            raise Exception(f"Error during API call: {e}")

    def __generate_qa(self, text_content, question_text=None):
        prompt = PROMPT_2.replace(
            "{{此处替换成你上一步生成的问题}}", question_text
        ).replace("{{此处替换成你的内容}}", text_content)
        messages = [
            {"role": "system", "content": prompt},
            {"role": "user", "content": "拼成问答对"},
        ]
        response = self.__generate_chat_completion(messages)
        return response.choices[0].message.content

    def __generate_question(self, text_content, more=False):
        content = "生成适合作为问答对的问题"
        if more:
            content = "尽可能多生成适合作为问答对的问题"
        prompt = PROMPT_1.replace("{{此处替换成你的内容}}", text_content)
        messages = [
            {"role": "system", "content": prompt},
            {"role": "user", "content": content},
        ]
        response = self.__generate_chat_completion(messages)
        return response.choices[0].message.content

    def __length_function(self, text: str) -> int:
        return len(self.enc.encode(text))

    def gen(self, file_name):
        knowledge_name = file_name.replace(".txt", "")
        text = minio_client.get_txt("books", file_name)

        splitter = RecursiveCharacterTextSplitter.from_language(
            Language.PYTHON,
            chunk_size=1500,
            chunk_overlap=200,
            length_function=self.__length_function,
        )
        splits = splitter.split_text(text)
        for txt in list(set(splits))[:1]:
            self.create_qa(txt)
        return len(splits)

    def create_qa(self, text_content):

        question_text = self.__generate_question(text_content=text_content, more=True)
        if not question_text:
            raise Exception("生成问题失败")
        qa_text = self.__generate_qa(
            text_content=text_content, question_text=question_text
        )
        if not qa_text:
            raise Exception("生成问答对失败")
        print(qa_text)
        return qa_text


if __name__ == "__main__":
    genqa_manager = GenqaManger()
    genqa_manager.gen("恶性犯罪深层心理探究.txt")
